"""
空调系统负荷预测数据预处理模块
包含数据清洗、特征工程、标准化等功能
"""

import pandas as pd
import numpy as np
from typing import Tuple, Dict, List, Optional
import datetime
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from sklearn.impute import SimpleImputer
import warnings
warnings.filterwarnings('ignore')


class HVACDataProcessor:
    """空调负荷数据预处理器"""
    
    def __init__(self, config: dict = None):
        """
        初始化数据处理器
        
        Args:
            config: 配置字典，包含处理参数
        """
        self.config = config or self._default_config()
        
        # 标准化器
        self.load_scaler = StandardScaler()
        self.numerical_scaler = StandardScaler()
        
        # 缺失值处理器
        self.load_imputer = SimpleImputer(strategy='median')
        self.numerical_imputer = SimpleImputer(strategy='mean')
        
        # 是否已拟合
        self.is_fitted = False
        
        # 特征列名
        self.numerical_features = [
            'outdoor_temp',      # 室外温度 (°C)
            'outdoor_humidity',  # 室外湿度 (%)
            'wind_speed',        # 风速 (m/s)
            'occupancy',         # 人员数量
            'equipment_status',  # 设备运行状态 (0-1)
            'electricity_price'  # 电价 (元/kWh)
        ]
        
    def _default_config(self) -> dict:
        """默认配置"""
        return {
            'seq_len': 168,          # 输入序列长度（一周）
            'pred_len': 24,          # 预测长度（一天）
            'outlier_threshold': 3,   # 异常值检测阈值（标准差倍数）
            'missing_threshold': 0.3, # 缺失值阈值
            'smooth_window': 3,       # 平滑窗口大小
        }
    
    def detect_outliers(self, data: np.ndarray, threshold: float = 3) -> np.ndarray:
        """
        检测异常值（基于Z-score）
        
        Args:
            data: 输入数据
            threshold: 异常值阈值
            
        Returns:
            is_outlier: 布尔数组，标识异常值位置
        """
        z_scores = np.abs((data - np.mean(data)) / np.std(data))
        return z_scores > threshold
    
    def smooth_data(self, data: np.ndarray, window: int = 3) -> np.ndarray:
        """
        数据平滑处理
        
        Args:
            data: 输入数据
            window: 滑动窗口大小
            
        Returns:
            smoothed_data: 平滑后的数据
        """
        if len(data) < window:
            return data
        
        smoothed = np.convolve(data, np.ones(window)/window, mode='same')
        # 处理边界效应
        smoothed[:window//2] = data[:window//2]
        smoothed[-(window//2):] = data[-(window//2):]
        
        return smoothed
    
    def extract_time_features(self, timestamps: pd.DatetimeIndex) -> np.ndarray:
        """
        提取时间特征
        
        Args:
            timestamps: 时间戳序列
            
        Returns:
            time_features: [n_samples, 4] 时间特征数组
        """
        features = np.zeros((len(timestamps), 4), dtype=np.int32)
        
        features[:, 0] = timestamps.hour                    # 小时 (0-23)
        features[:, 1] = timestamps.dayofweek              # 星期 (0-6)
        features[:, 2] = timestamps.month - 1              # 月份 (0-11)
        features[:, 3] = (timestamps.month - 1) // 3       # 季节 (0-3)
        
        return features
    
    def add_lag_features(self, data: np.ndarray, lags: List[int] = None) -> np.ndarray:
        """
        添加滞后特征
        
        Args:
            data: 原始数据 [n_samples, n_features]
            lags: 滞后期列表
            
        Returns:
            data_with_lags: 包含滞后特征的数据
        """
        if lags is None:
            lags = [1, 2, 3, 24, 168]  # 1小时、2小时、3小时、1天、1周前
        
        n_samples, n_features = data.shape
        lag_features = []
        
        for lag in lags:
            if lag < n_samples:
                # 创建滞后特征
                lag_data = np.zeros_like(data)
                lag_data[lag:] = data[:-lag]
                lag_features.append(lag_data)
        
        if lag_features:
            return np.concatenate([data] + lag_features, axis=1)
        else:
            return data
    
    def calculate_statistical_features(self, data: np.ndarray, window: int = 24) -> np.ndarray:
        """
        计算统计特征（滚动窗口）
        
        Args:
            data: 输入数据
            window: 滚动窗口大小
            
        Returns:
            statistical_features: 统计特征
        """
        features = []
        
        for i in range(len(data)):
            start_idx = max(0, i - window + 1)
            window_data = data[start_idx:i+1]
            
            if len(window_data) > 0:
                features.append([
                    np.mean(window_data),     # 均值
                    np.std(window_data),      # 标准差
                    np.min(window_data),      # 最小值
                    np.max(window_data),      # 最大值
                ])
            else:
                features.append([0, 0, 0, 0])
        
        return np.array(features)
    
    def create_holiday_features(self, timestamps: pd.DatetimeIndex) -> np.ndarray:
        """
        创建节假日特征
        
        Args:
            timestamps: 时间戳序列
            
        Returns:
            holiday_features: 节假日特征
        """
        # 简化版节假日检测（可以根据需要扩展）
        features = np.zeros((len(timestamps), 3))
        
        # 周末特征
        features[:, 0] = (timestamps.dayofweek >= 5).astype(int)
        
        # 工作时间特征（8-18点）
        features[:, 1] = ((timestamps.hour >= 8) & (timestamps.hour <= 18)).astype(int)
        
        # 高峰时间特征（10-16点）
        features[:, 2] = ((timestamps.hour >= 10) & (timestamps.hour <= 16)).astype(int)
        
        return features
    
    def fit(self, load_data: np.ndarray, numerical_data: np.ndarray, 
            timestamps: pd.DatetimeIndex):
        """
        拟合预处理器
        
        Args:
            load_data: 负荷数据
            numerical_data: 数值特征数据
            timestamps: 时间戳
        """
        print("正在拟合数据预处理器...")
        
        # 处理缺失值
        load_data_clean = self.load_imputer.fit_transform(load_data.reshape(-1, 1)).flatten()
        numerical_data_clean = self.numerical_imputer.fit_transform(numerical_data)
        
        # 异常值检测和处理
        load_outliers = self.detect_outliers(load_data_clean, self.config['outlier_threshold'])
        print(f"检测到 {np.sum(load_outliers)} 个负荷异常值")
        
        # 用中位数替换异常值
        load_data_clean[load_outliers] = np.median(load_data_clean)
        
        # 拟合标准化器
        self.load_scaler.fit(load_data_clean.reshape(-1, 1))
        self.numerical_scaler.fit(numerical_data_clean)
        
        self.is_fitted = True
        print("数据预处理器拟合完成")
    
    def transform(self, load_data: np.ndarray, numerical_data: np.ndarray, 
                 timestamps: pd.DatetimeIndex) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        转换数据
        
        Args:
            load_data: 负荷数据
            numerical_data: 数值特征数据
            timestamps: 时间戳
            
        Returns:
            processed_load: 处理后的负荷数据
            processed_numerical: 处理后的数值特征
            time_features: 时间特征
        """
        if not self.is_fitted:
            raise ValueError("预处理器尚未拟合，请先调用 fit() 方法")
        
        # 处理缺失值
        load_data_clean = self.load_imputer.transform(load_data.reshape(-1, 1)).flatten()
        numerical_data_clean = self.numerical_imputer.transform(numerical_data)
        
        # 异常值处理
        load_outliers = self.detect_outliers(load_data_clean, self.config['outlier_threshold'])
        load_data_clean[load_outliers] = np.median(load_data_clean)
        
        # 数据平滑
        if self.config['smooth_window'] > 1:
            load_data_clean = self.smooth_data(load_data_clean, self.config['smooth_window'])
        
        # 标准化
        processed_load = self.load_scaler.transform(load_data_clean.reshape(-1, 1)).flatten()
        processed_numerical = self.numerical_scaler.transform(numerical_data_clean)
        
        # 提取时间特征
        time_features = self.extract_time_features(timestamps)
        
        # 添加额外特征
        holiday_features = self.create_holiday_features(timestamps)
        statistical_features = self.calculate_statistical_features(load_data_clean)
        
        # 合并数值特征
        enhanced_numerical = np.concatenate([
            processed_numerical,
            holiday_features,
            statistical_features
        ], axis=1)
        
        return processed_load, enhanced_numerical, time_features
    
    def inverse_transform_load(self, scaled_load: np.ndarray) -> np.ndarray:
        """
        反标准化负荷数据
        
        Args:
            scaled_load: 标准化后的负荷数据
            
        Returns:
            original_load: 原始尺度的负荷数据
        """
        return self.load_scaler.inverse_transform(scaled_load.reshape(-1, 1)).flatten()
    
    def create_sequences(self, load_data: np.ndarray, numerical_data: np.ndarray, 
                        time_features: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """
        创建训练序列
        
        Args:
            load_data: 负荷数据
            numerical_data: 数值特征数据
            time_features: 时间特征
            
        Returns:
            X_load: 输入负荷序列
            X_numerical: 输入数值特征序列
            X_time: 输入时间特征序列
            y: 目标负荷序列
        """
        seq_len = self.config['seq_len']
        pred_len = self.config['pred_len']
        
        n_samples = len(load_data) - seq_len - pred_len + 1
        
        X_load = np.zeros((n_samples, seq_len, 1))
        X_numerical = np.zeros((n_samples, seq_len, numerical_data.shape[1]))
        X_time = np.zeros((n_samples, seq_len, time_features.shape[1]))
        y = np.zeros((n_samples, pred_len, 1))
        
        for i in range(n_samples):
            X_load[i] = load_data[i:i+seq_len].reshape(-1, 1)
            X_numerical[i] = numerical_data[i:i+seq_len]
            X_time[i] = time_features[i:i+seq_len]
            y[i] = load_data[i+seq_len:i+seq_len+pred_len].reshape(-1, 1)
        
        return X_load, X_numerical, X_time, y
    
    def get_feature_info(self) -> Dict:
        """获取特征信息"""
        return {
            'numerical_features': self.numerical_features,
            'time_features': ['hour', 'day_of_week', 'month', 'season'],
            'enhanced_features': ['weekend', 'work_hours', 'peak_hours', 
                                'load_mean', 'load_std', 'load_min', 'load_max'],
            'total_numerical_features': 6 + 3 + 4,  # 原始特征 + 节假日特征 + 统计特征
        }


def load_sample_data() -> Tuple[np.ndarray, np.ndarray, pd.DatetimeIndex]:
    """
    生成示例数据（实际使用时替换为真实数据加载函数）
    
    Returns:
        load_data: 负荷数据
        numerical_data: 数值特征数据
        timestamps: 时间戳
    """
    # 生成一个月的小时级数据
    n_hours = 24 * 30  # 30天
    timestamps = pd.date_range('2024-01-01', periods=n_hours, freq='H')
    
    # 生成模拟负荷数据（带有日周期和噪声）
    hours = np.arange(n_hours) % 24
    days = np.arange(n_hours) // 24
    
    # 基础负荷模式（白天高，夜晚低）
    base_load = 100 + 50 * np.sin(2 * np.pi * hours / 24 - np.pi/2)
    
    # 周循环（工作日vs周末）
    week_pattern = 1 + 0.2 * np.sin(2 * np.pi * days / 7)
    
    # 添加噪声
    noise = np.random.normal(0, 5, n_hours)
    
    load_data = base_load * week_pattern + noise
    
    # 生成数值特征
    numerical_data = np.random.randn(n_hours, 6)
    numerical_data[:, 0] = 20 + 10 * np.sin(2 * np.pi * days / 365) + np.random.normal(0, 2, n_hours)  # 温度
    numerical_data[:, 1] = 50 + 20 * np.random.rand(n_hours)  # 湿度
    numerical_data[:, 2] = 2 + 3 * np.random.rand(n_hours)    # 风速
    numerical_data[:, 3] = 50 + 30 * (hours > 8) * (hours < 18)  # 人员数量
    numerical_data[:, 4] = np.random.choice([0, 1], n_hours, p=[0.1, 0.9])  # 设备状态
    numerical_data[:, 5] = 0.5 + 0.3 * np.random.rand(n_hours)  # 电价
    
    return load_data, numerical_data, timestamps


if __name__ == "__main__":
    # 示例使用
    print("=== 空调负荷数据预处理示例 ===")
    
    # 加载示例数据
    load_data, numerical_data, timestamps = load_sample_data()
    print(f"数据形状: 负荷 {load_data.shape}, 数值特征 {numerical_data.shape}, 时间戳 {len(timestamps)}")
    
    # 创建数据处理器
    processor = HVACDataProcessor()
    
    # 拟合处理器
    processor.fit(load_data, numerical_data, timestamps)
    
    # 转换数据
    processed_load, processed_numerical, time_features = processor.transform(
        load_data, numerical_data, timestamps
    )
    
    print(f"处理后数据形状:")
    print(f"  负荷: {processed_load.shape}")
    print(f"  数值特征: {processed_numerical.shape}")
    print(f"  时间特征: {time_features.shape}")
    
    # 创建训练序列
    X_load, X_numerical, X_time, y = processor.create_sequences(
        processed_load, processed_numerical, time_features
    )
    
    print(f"训练序列形状:")
    print(f"  X_load: {X_load.shape}")
    print(f"  X_numerical: {X_numerical.shape}")
    print(f"  X_time: {X_time.shape}")
    print(f"  y: {y.shape}")
    
    # 显示特征信息
    feature_info = processor.get_feature_info()
    print(f"特征信息: {feature_info}")
